import numpy as np
from numpy.ma import cos
import matplotlib.pyplot as plt
from matplotlib import cm
from mpl_toolkits.mplot3d import Axes3D
import math
import matplotlib

matplotlib.use('TkAgg')


class GeneticAlgorithmOptimizer:
    def __init__(self):
        # 算法参数
        self.DNA_SIZE = 24  # 编码长度
        self.POP_SIZE = 200  # 种群大小
        self.CROSS_RATE = 0.5  # 交叉率
        self.MUTA_RATE = 0.015  # 变异率
        self.Iterations = 100  # 迭代次数
        self.X_BOUND = [1, 2]  # X区间
        self.Y_BOUND = [1, 2]  # Y区间

        # 初始化种群
        self.population = np.random.randint(2, size=(self.POP_SIZE, self.DNA_SIZE * 2))

        # 记录迭代过程
        self.best_fitness_history = []
        self.average_fitness_history = []
        self.best_individual_history = []

        # 创建可视化界面
        self.fig = plt.figure(figsize=(16, 8))

        # 3D可视化设置
        self.ax_3d = self.fig.add_subplot(121, projection='3d')
        self.setup_3d_visualization()

        # 2D折线图设置
        self.ax_2d = self.fig.add_subplot(122)
        self.setup_2d_visualization()

        # 开启交互模式
        plt.ion()

    def fitness_function(self, x, y):
        """适应度函数 - 计算目标函数值"""
        return 20 + x ** 2 + y ** 2 - 10 * (cos(2 * math.pi * x) + cos(2 * math.pi * y))

    def decode_genome(self, population):
        """解码DNA序列为实际参数值"""
        x_pop = population[:, 1::2]  # 奇数列表示X
        y_pop = population[:, ::2]  # 偶数列表示y

        # 将二进制转换为实数
        x = x_pop.dot(2 ** np.arange(self.DNA_SIZE)[::-1]) / float(2 ** self.DNA_SIZE - 1) * \
            (self.X_BOUND[1] - self.X_BOUND[0]) + self.X_BOUND[0]
        y = y_pop.dot(2 ** np.arange(self.DNA_SIZE)[::-1]) / float(2 ** self.DNA_SIZE - 1) * \
            (self.Y_BOUND[1] - self.Y_BOUND[0]) + self.Y_BOUND[0]

        return x, y

    def calculate_fitness(self, population):
        """计算种群适应度"""
        x, y = self.decode_genome(population)
        values = self.fitness_function(x, y)
        # 转换为最小化问题，防止适应度出现负数
        return -(values - np.max(values)) + 0.0001

    def selection(self, population, fitness):
        """选择操作 - 轮盘赌选择"""
        idx = np.random.choice(
            np.arange(self.POP_SIZE),
            size=self.POP_SIZE,
            replace=True,
            p=fitness / fitness.sum()
        )
        return population[idx]

    def crossover_and_mutation(self, population):
        """交叉和变异操作"""
        new_population = []
        for individual in population:
            # 交叉操作
            if np.random.rand() < self.CROSS_RATE:
                mate_idx = np.random.randint(self.POP_SIZE)
                crossover_point1 = np.random.randint(0, self.DNA_SIZE * 2 - 1)
                crossover_point2 = np.random.randint(crossover_point1, self.DNA_SIZE * 2)
                individual[crossover_point1:crossover_point2] = \
                    population[mate_idx][crossover_point1:crossover_point2]

            # 变异操作
            self.mutate(individual)

            new_population.append(individual)

        return np.array(new_population)

    def mutate(self, individual):
        """变异操作"""
        if np.random.rand() < self.MUTA_RATE:
            mutate_point = np.random.randint(0, self.DNA_SIZE)
            individual[mutate_point] = individual[mutate_point] ^ 1

    def setup_3d_visualization(self):
        """初始化3D可视化设置"""
        # 绘制目标函数曲面
        X = np.linspace(*self.X_BOUND, 100)
        Y = np.linspace(*self.Y_BOUND, 100)
        X, Y = np.meshgrid(X, Y)
        Z = self.fitness_function(X, Y)

        self.ax_3d.plot_surface(X, Y, Z, rstride=1, cstride=1,
                                cmap=cm.coolwarm, alpha=0.6)
        self.ax_3d.set_zlim(0, 40)
        self.ax_3d.set_xlabel('x')
        self.ax_3d.set_ylabel('y')
        self.ax_3d.set_zlabel('z')
        self.ax_3d.set_title('3D Function Visualization with Population')

        # 显示初始种群
        x, y = self.decode_genome(self.population)
        self.scatter_3d = self.ax_3d.scatter(
            x, y, self.fitness_function(x, y),
            c='red', marker='o', s=30, alpha=0.8
        )

    def setup_2d_visualization(self):
        """初始化2D折线图设置"""
        self.ax_2d.set_xlim(0, self.Iterations)
        self.ax_2d.set_ylim(0, 40)  # 根据函数值范围设置
        self.ax_2d.set_xlabel('Generation')
        self.ax_2d.set_ylabel('Fitness Value')
        self.ax_2d.set_title('Fitness Trend Over Generations')
        self.ax_2d.grid(True)

        # 初始化线条
        self.best_line, = self.ax_2d.plot([], [], 'r-', label='Best Fitness')
        self.avg_line, = self.ax_2d.plot([], [], 'b-', label='Average Fitness')
        self.ax_2d.legend()

    def update_visualizations(self, iteration):
        """更新所有可视化元素"""
        x, y = self.decode_genome(self.population)
        fitness = self.calculate_fitness(self.population)

        # 更新3D可视化
        self.scatter_3d.remove()
        self.scatter_3d = self.ax_3d.scatter(
            x, y, self.fitness_function(x, y),
            c='red', marker='o', s=30, alpha=0.8
        )
        self.ax_3d.set_title(f'Population at Generation {iteration + 1}')

        # 记录当前代的统计数据
        self.best_fitness_history.append(np.max(fitness))
        self.average_fitness_history.append(np.mean(fitness))

        # 更新2D折线图
        generations = range(iteration + 1)
        self.best_line.set_data(generations, self.best_fitness_history)
        self.avg_line.set_data(generations, self.average_fitness_history)

        # 调整Y轴范围以适应数据
        current_max = max(self.best_fitness_history)
        self.ax_2d.set_ylim(0, current_max * 1.2)

        # 重绘图形
        plt.draw()
        plt.pause(0.05)

    def print_results(self):
        """打印最终结果"""
        fitness = self.calculate_fitness(self.population)
        best_idx = np.argmin(fitness)
        x, y = self.decode_genome(self.population)

        print("\n========== Optimization Results ==========")
        print(f"Best fitness: {fitness[best_idx]}")
        print(f"Best genome: {self.population[best_idx]}")
        print(f"Best (x, y): ({x[best_idx]}, {y[best_idx]})")
        print(f"Minimum F(x,y): {self.fitness_function(x[best_idx], y[best_idx])}")

    def run(self):
        """运行遗传算法优化"""
        print('========== Genetic Algorithm Optimization ==========')

        for iteration in range(self.Iterations):
            print(f"Iteration {iteration + 1}/{self.Iterations}")

            # 计算适应度
            fitness = self.calculate_fitness(self.population)

            # 记录最佳结果
            best_idx = np.argmin(fitness)
            x, y = self.decode_genome(self.population)
            self.best_individual_history.append((x[best_idx], y[best_idx]))

            # 更新可视化
            self.update_visualizations(iteration)

            # 遗传操作
            self.population = self.crossover_and_mutation(self.population)
            fitness = self.calculate_fitness(self.population)
            self.population = self.selection(self.population, fitness)

        # 显示最终结果
        self.print_results()

        # 关闭交互模式并保持图形显示
        plt.ioff()
        plt.show()


if __name__ == "__main__":
    optimizer = GeneticAlgorithmOptimizer()
    optimizer.run()